# core
from multiprocessing.dummy import Array
from tkinter import OFF
import draw

from importlib.metadata import entry_points
import string
from typing import Set
import idc
import idaapi
import ida_kernwin
import ida_funcs
import idautils
import ida_range
from example.ida.utils import guess_machine,expr2colorstr
from future.utils import viewitems, viewvalues

from miasm.arch.x86.regs import *
from miasm.core.bin_stream_ida import bin_stream_ida
from miasm.expression.simplifications import expr_simp
from miasm.ir.ir import IRBlock, AssignBlock
from miasm.ir.symbexec import SymbolicExecutionEngine
from miasm.analysis.data_flow import DeadRemoval, merge_blocks, remove_empty_assignblks
from miasm.analysis.cst_propag import propagate_cst_expr
from miasm.analysis.data_flow import load_from_int
from miasm.analysis.data_flow import ReachingDefinitions,DiGraphDefUse
from miasm.expression.expression import ExprLoc, ExprInt, ExprOp, ExprAssign,ExprCond, ExprId, ExprMem,ExprSlice
from miasm.analysis.simplifier import IRCFGSimplifierCommon, IRCFGSimplifierSSA
from miasm.core.locationdb import LocationDB
from miasm.core.asmblock import AsmCFG,AsmBlock,AsmConstraint
from miasm.ir.translators.translator import Translator
from miasm.ir.ir import IRCFG

from mod_utils import *
# from BaseSolver import BaseSolver
import BaseSolver
from scanner import *
import draw
import ollvm
import mylib
import scanner
import igraph as ig
from collections import Counter
# from mylib import *
import sys
from PyQt5.QtWidgets import (QApplication,QWidget,QVBoxLayout,
                             QListView,QMessageBox,QAbstractItemView,
                             QGridLayout,QPushButton,QLabel,QCheckBox,QTextEdit)
from PyQt5.QtCore import QStringListModel
from PyQt5 import QtCore, QtGui, QtWidgets
import json

class MyPluginFormClass(ida_kernwin.PluginForm):
    log_format=None
    def OnCreate(self, form):
        """
        Called when the widget is created
        """

        # Get parent widget
        self.parent = self.FormToPyQtWidget(form)
        self.PopulateForm()

    def PopulateForm(self):
        # Create layout
        layout = QtWidgets.QVBoxLayout()
        grid = QGridLayout()
        grid.setSpacing(5)
        self.analysis_btn=QPushButton('Analysis')
        self.pre_analysis_btn=QPushButton('Pre Analysis')
        self.graph_info_btn=QPushButton('Graph Info')
        self.log_btn=QPushButton('Log')
        hello_label=QLabel("Code With <font color=red>Miasm</font> from <font color=blue>IDAPython</font>")
        rule_label=QLabel('<font color=red>Rule List</font>')
        #des_label=QLabel('Yet another great obfuscator by IDA python')
        des_label=QLabel('-----------------------------------------')
        blank_label=QLabel('-----------------------------------------')
        self.pattern_match_checkbox=QCheckBox('Pattern Matching')
        self.simplify_result_checkbox=QCheckBox('Simplify Result')
        self.log_edit = QTextEdit()
        # 默认勾选
        # pattern_match_checkbox.toggle()
        self.log_edit.setFocusPolicy(QtCore.Qt.NoFocus)
        
        #实例化列表视图
        self.listview=QListView()
        self.listview.setEditTriggers(QAbstractItemView.NoEditTriggers)
        #实例化列表模型，添加数据
        slm=QStringListModel()
        self.qList=['OLLVM-fla','OLLVM-bcf','OLLVM-sub','emotet','maze']

        #设置模型列表视图，加载数据列表
        slm.setStringList(self.qList)

        #设置列表视图的模型
        self.listview.setModel(slm)

        #单击触发自定义的槽函数
        # listview.doubleclicked.connect(self.clicked)
        self.analysis_btn.clicked.connect(self.onClickeAnalysis)
        self.pre_analysis_btn.clicked.connect(self.onClickPreAnalysis)
        self.graph_info_btn.clicked.connect(self.onClickGraphInfo)
        self.log_btn.clicked.connect(self.onClickLog)
        #设置窗口布局，加载控件
        grid.addWidget(hello_label,1,0,1,1)
        # grid.addWidget(blank_label,1,1)
        # grid.addWidget(des_label,1,2)
        # grid.addWidget(blank_label,1,3)
        grid.addWidget(self.pattern_match_checkbox,2,0)
        grid.addWidget(self.simplify_result_checkbox,3,0)
        grid.addWidget(rule_label,4,0)
        grid.addWidget(self.pre_analysis_btn,4,4)
        grid.addWidget(self.graph_info_btn,4,5)
        grid.addWidget(self.log_btn,4,6)
        
        grid.addWidget(self.listview,5,0,4,2)
        grid.addWidget(self.log_edit,5,4,4,4)
        grid.addWidget(self.analysis_btn,10,0,1,1)
        # grid.addWidget(blank_label,8,3,1,1)
        
        self.parent.setLayout(grid)

    def onClickeAnalysis(self):
        # index=qModelIndex.row()
        format=FormatLines()
        index = self.listview.currentIndex().row()
        RuleSolver= core(index,self.pattern_match_checkbox.isChecked())
        format._add_line(("[x]Analysis Completed,Rule Name [%s]") % (RuleSolver.rule_name))
        str_add_blk,str_del_blk,str_add_edge,str_del_edge=RuleSolver.format_action()
        if len(str_add_blk)>0:
            json_add_blk=json.dumps(str_add_blk)
            format._add_line('[added_block]')
            format._add_line('\t'+json_add_blk)
        if len(str_del_blk)>0:
            json_del_blk=json.dumps(str_del_blk)
            format._add_line('[deleted_block]')
            format._add_line('\t'+json_del_blk)
            pass
        if len(str_add_edge)>0:
            json_add_edge=json.dumps(str_add_edge)
            format._add_line('[added_edge]')
            format._add_line('\t'+json_add_edge)
            pass
        if len(str_del_edge)>0:
            json_del_edge=json.dumps(str_del_edge)
            format._add_line('[deleted_edge]')
            format._add_line('\t'+json_del_edge)
            pass
        self.log_format=format
        self.log_edit.setPlainText(format.str())
        return
    def onClickPreAnalysis(self):
        ida_obj=pre_analysis()
        format=FormatLines()
        format.add_prompt('Pre Analysis')
        for key in ida_obj:
            value=ida_obj[key]
            if isinstance(value,int):
                if value>0xFFFF:
                    value=hex(value)
                else:
                    value=str(value)
            tmp_value=[]
            if isinstance(value,list):
                for vv in value:
                    tmp_value.append((hex(vv[0]),hex(vv[1])))
                value=tmp_value
            format.add_key_value(key,value)
                    
        self.log_edit.setPlainText(format.str())
        return

    def onClickGraphInfo(self):
        format=FormatLines()
        format.add_prompt('Graph info')
        load_obj=load_database()
        asmcfg=load_obj['asmcfg']
        g=graph_analysis(asmcfg)
        graph_obj=g.degree()
        blk_count=graph_obj['blk_count']
        in_cnt=graph_obj['in_cnt']
        out_cnt=graph_obj['out_cnt']
        io_cnt=graph_obj['ipo_cnt']
        format.add_key_value('block count',blk_count)
        format.add_dict('in_degree',in_cnt)
        format.add_dict('out_degree',out_cnt)
        format.add_dict('io_degree',io_cnt)
        
        
        self.log_edit.setPlainText(format.str())
        
        return

    def onClickLog(self):
        self.log_edit.setPlainText(self.log_format.str())
        return
    def OnClose(self, form):
        """
        Called when the widget is closed
        """
        pass


# 每个node即为一个asmblock
# 每个
class CFGraph:
    def __init__(self) -> None:
        self.g=ig.Graph(directed=True)
        pass
    #@offset:汇编块的第一块的地址
    #@len:汇编指令的条数
    def add_node(self,offset,length):
        new_id=len(self.g.vs)   #新的块id
        self.g.add_vertex(None,id=new_id,offset=offset,length=length)
        return new_id
    def add_edge(self,src,dst,c_t):        
        self.g.add_edge(src,dst,src=src,dst=dst,type=c_t)

    def get_id_by_offset(self,offset):
        ret=self.g.vs.select(offset=offset)['id']
        if len(ret)>1:
            assert()
        return ret[0]
    def get_node_attr(self,n_id):
        p_node=self.g.vs[n_id]
        return {
            'id':p_node['id'],
            'offset':p_node['offset'],
            'length':p_node['length']
        }
    def find_edge(self,src):
        p_edge=self.g.es.select(_source=src)
        # for edge in p_edge:
        #     print(edge.source,edge.target)
        return
    def degree(self):
        alls=self.g.degree()
        ins=self.g.degree(mode='in')
        outs=self.g.degree(mode='out')
        ipos=[0]*len(alls)
        for index in range(0,len(alls)):
            if ins[index]!=0 and outs[index]!=0:
                ipos[index]=ins[index]/outs[index]
 
        ins_cnt=Counter(ins)
        outs_cnt=Counter(outs)
        ipo_cnt=Counter(ipos)    
        return {
            'blk_count':len(alls),
            'all':alls,
            'in':ins,
            'out':outs,
            'ipo':ipos,
            'in_cnt':dict(ins_cnt),
            'out_cnt':dict(outs_cnt),
            'ipo_cnt':dict(ipo_cnt)
        }
    def show(self):
        print(self.g)
# 一个规范的输出格式
class FormatLines:
    def __init__(self) -> None:
        self.s=''
        pass
    def _add_line(self,line):
        self.s+=line
        self.s+='\n'
    def add_prompt(self,info):
        prompt=("[x]%s") % (info)
        self._add_line(prompt)
    def add_key_value(self,key,value):
        key_value=("[%s]\t%s") % (key,value)
        self._add_line(key_value)
    def add_dict(self,key,obj):
        key_str=("[%s]") % (key)
        self._add_line(key_str)
        for key in obj:
            obj_str=("\t[%s]\t%s") % (str(key),str(obj[key]))
            self._add_line(obj_str)
    
    def str(self):
        return self.s
       
def pre_analysis():
    ida_obj={
        'file name':idc.get_root_filename(),
        'screen_ea':None,
        'is_func':False
    }
    screen_ea=idc.get_screen_ea()
    ida_obj['screen_ea']=screen_ea
    func_now=ida_funcs.get_func(screen_ea)
    if func_now is None:
        return ida_obj
    ida_obj['is_func']=True
    ida_obj['func_ea']=func_now.start_ea
    ida_obj['func_size']=ida_funcs.calc_func_size(func_now)
    ida_obj['func_bitness']=ida_funcs.get_func_bitness(func_now)
    ida_obj['func_bits']=ida_funcs.get_func_bits(func_now)
    ida_obj['func_bytes']=ida_funcs.get_func_bytes(func_now)
    ida_obj['func_name']=ida_funcs.get_func_name(func_now.start_ea)
    func_ranges=ida_range.rangeset_t()
    ida_funcs.get_func_ranges(func_ranges,func_now)
    func_range_arr=[]
    for item in func_ranges:
        func_range_arr.append((item.start_ea,item.end_ea))
    ida_obj['func_range']=func_range_arr
    ida_obj['is_entry']=ida_funcs.is_func_entry(func_now)
    ida_obj['is_tail']=ida_funcs.is_func_tail(func_now)
    ida_obj['is_locked']=ida_funcs.is_func_locked(func_now)
    ida_obj['is_visible']=ida_funcs.is_visible_func(func_now)
    ida_obj['is_sp']=func_now.analyzed_sp()
    ida_obj['is_return']=func_now.does_return()
    ida_obj['is_far']=func_now.is_far()
    ida_obj['need_prolog']=func_now.need_prolog_analysis()
    return ida_obj
    
def graph_analysis(asmcfg:AsmCFG):
    # init graph
    edges=[]
    g=CFGraph()
    for asmblk in asmcfg.blocks:
        blk_off=asmblk.get_offsets()[0]
        blk_len=len(asmblk.get_offsets())
        g.add_node(blk_off,blk_len)
        # print(hex(blk_off),blk_len,asmblk.bto)
        for bto in asmblk.bto:
            bto_off=asmcfg.loc_db.loc_key_to_offset(bto.loc_key)
            edges.append((blk_off,bto_off,bto.c_t))
    # print(edges)
    for edge in edges:
        src_id=g.get_id_by_offset(edge[0])
        dst_id=g.get_id_by_offset(edge[1])
        # print(src_id,'->',dst_id)
        g.add_edge(src_id,dst_id,edge[2])
    return g


def load_database():
    print('[x]start analysis')
    
    # get function
    screen_ea=idc.get_screen_ea()
    func_now=ida_funcs.get_func(screen_ea)
    if func_now is None:
        now_addr=screen_ea
    else:
        now_addr=func_now.start_ea
        
        # print('refqt',func_now.refqty)
        # print(func_now.referers.data)

    
    # arch info
    machine=guess_machine(addr=now_addr)
    dis_engine=machine.dis_engine
    
    # 从idb获得二进制数据
    main_bs=bin_stream_ida()
    main_loc_db=LocationDB()
    main_mdis=dis_engine(main_bs,loc_db=main_loc_db)
        
    for addr,name in idautils.Names():
        if name is None:
            continue
        if(main_loc_db.get_offset_location(addr) or main_loc_db.get_name_location(name)):
            continue
        main_loc_db.add_location(name,addr)
    # main_AsmCFG
    main_asmcfg=main_mdis.dis_multiblock(now_addr)
    # 设置反汇编的边界，在jmp处停止
    # main_mdis.dis_block_callback=stop_on_jmp
    
    # 得到入口loc_key
    entry_points=set([main_loc_db.get_offset_location(now_addr)])
    
    # 流程图的头节点
    main_head=list(entry_points)[0]
    res_obj={
        "fname":None,
        "func_now":func_now,
        "addr":now_addr,
        "machine":machine,
        "bs":main_bs,
        "loc_db":main_loc_db,
        "dis":main_mdis,
        "asmcfg":main_asmcfg,
        "entry":entry_points,
        "head":main_head,
    }
    return res_obj
import traceback
def core(rule_index,is_fattern):
    # load database
    load_obj=load_database()
    asmcfg=load_obj['asmcfg']
    head=load_obj['head']
    
    mdis=load_obj['dis']
    loc_db=load_obj['loc_db']
    func_now=load_obj['func_now']
    machine=load_obj['machine']
    addr=load_obj['addr']
    RuleSolver=None
    graph_analysis(asmcfg)
    # return
    if is_fattern:
        print('Fattern matcing...')
        scanner.pattern_matching(asmcfg)
    try:
        if rule_index==0:
            # ollvm.test_func(1)
            RuleSolver=ollvm.OllvmFlaSolver()
            RuleSolver.Solve(load_obj)
            # lift_obj=make_lift(machine,loc_db,asmcfg)
            # ircfg=lift_obj['ircfg']
            # lifter=lift_obj['lifter']
            # ollvm.fla_ollvm(asmcfg,head,ircfg,lifter,mdis,loc_db)
        elif rule_index==1:
            RuleSolver=ollvm.OllvmBcfSolver()
            RuleSolver.Solve(load_obj)
            # lift_obj=make_lift(machine,loc_db,asmcfg)
            # ircfg=lift_obj['ircfg']
            # lifter=lift_obj['lifter']
            # ollvm.bcf_ollvm(func_now,machine,asmcfg,head,ircfg,lifter,mdis,loc_db)
        elif rule_index==2:
            scanner.pattern_matching(asmcfg)
            draw.draw_graph(asmcfg,'fix_asmcfg')
        elif rule_index==3:
            RuleSolver=emotet.EmotetSolver()
            RuleSolver.Solve(load_obj)
            # emotet.fla_emotet(func_now,machine,asmcfg,head,ircfg,lifter,mdis,loc_db)
        elif rule_index==4:
            RuleSolver=maze.MazeSolver()
            RuleSolver.Solve(load_obj)
        
        
    except Exception as r:
        print('Exception',traceback.format_exc())
    return RuleSolver
    lift_obj=make_lift(machine,loc_db,asmcfg)
    ircfg=lift_obj['ircfg']
    lifter=lift_obj['lifter']
    
    # head_blk=asmcfg.loc_key_to_block(head)
    # head_off=head_blk.get_offsets()[0]
    # main_asmcfg=asmcfg
    # main_ircfg=ircfg
    ollvm.fla_ollvm(asmcfg,head,ircfg,lifter,mdis,loc_db)
    # ollvm.bcf_ollvm(func_now,machine,asmcfg,head,ircfg,lifter,mdis,loc_db)
    # ollvm.sub_ollvm(func_now,machine,asmcfg,head,ircfg,lifter,mdis,loc_db)
    # emotet

    # emotet.fla_emotet(func_now,machine,asmcfg,head,ircfg,lifter,mdis,loc_db)
    # maze
    # maze.sub_maze(addr,machine,asmcfg,head,mdis,loc_db)
    return

from importlib import reload
import emotet
import maze
if __name__=='__main__':
    reload(ollvm)
    reload(emotet)
    reload(maze)
    reload(mylib)
    reload(scanner)
    reload(draw)
    reload(BaseSolver)
    # pre_analysis()
    # core(1,False)
    
    plg = MyPluginFormClass()
    plg.Show("IDA Binary Obfuscator")
    
    
    # sys.exit(app.exec_())
    
    
    
